from stable_baselines3 import DDPG
from stable_baselines3.common.noise import NormalActionNoise
from init_control import  UAVParameters
from DDPG_uav_dynamics import UAVEnv
import numpy as np

params = UAVParameters()
waypoints = [
    np.array([0, 0, 10]),
    np.array([10, 0, 10]),
    np.array([10, 10, 10]),
    np.array([0, 10, 10]),
    np.array([0, 0, 10])
]  # Example waypoints

env = UAVEnv(params, waypoints)

# Add some noise to the actions for exploration
n_actions = env.action_space.shape[-1]
action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))

# Create the DDPG model
model = DDPG('MlpPolicy', env, action_noise=action_noise, verbose=1)

# Train the model
model.learn(total_timesteps=100000)

# Save the model
model.save("ddpg_uav")

# Load the model
model = DDPG.load("ddpg_uav")

# Evaluate the trained agent
obs = env.reset()
for _ in range(1000):
    action, _states = model.predict(obs)
    obs, rewards, dones, info = env.step(action)
    if dones:
        break

# Simulate and visualize UAV trajectory
positions = []

obs = env.reset()
for _ in range(1000):
    action, _states = model.predict(obs)
    obs, rewards, dones, info = env.step(action)
    positions.append(obs[:3])  # Record position (x, y, z)
    if dones:
        break

positions = np.array(positions)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot(positions[:, 0], positions[:, 1], positions[:, 2])
ax.set_xlabel('X Position')
ax.set_ylabel('Y Position')
ax.set_zlabel('Z Position')
ax.set_title('UAV Trajectory')
plt.show()
